#!/usr/bin/env python
#-*- coding: UTF-8 -*-
# schif.py: task manager interface functions for scheduler
#
# Copyright (C) 2008-2010 Hily Jiang <hilyjiang@gmail.com>
#
# This file is part of dtask.
#
# This software is provided 'as-is', without any express or implied
# warranty.  In no event will the authors be held liable for any damages
# arising from the use of this software.
#
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not
#    claim that you wrote the original software. If you use this software
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
#    misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.

"""
Task Manager interface functions for Scheduler.

Scheduler <-> Task Manager Protocol:

* New Connection
  (S -> T) Send data format:    PASSPORT schName schKey\n

* Get Task
  (S -> T) Send data format:    GETTASK\n
  (T -> S) Return data format:  TASK taskId taskType taskParams\n

* Write Result
  (S -> T) Send data format:    RESULT taskId result\n

* Update Status
  (S -> T) Send data format:    STATUS taskId statusType statusInfo\n
                                statuType:
                                  nodm: node manager's status
                                  nod:  task client's status
                                statuInfo:
                                  nodm:
                                    nodmip nodmport nodmast nodmrtt
                                  nod:
                                    nodname nodip nodport nodast nodrtt
"""

import sys
import threading

from utils import CommandBuffer


CFG_SCH_BUFSIZ = 1024


class SchIfBuffer(CommandBuffer):
    """Scheduler buffer to receive command string and scheduler info."""

    def __init__(self):
        CommandBuffer.__init__(self)
        self.schName = ''
        self.schKey = ''
        self.hasInfo = False

    def setInfo(self, schName='', schKey=''):
        """Set scheduler info."""
        if self.hasInfo:
            return False
        self.schName = schName
        self.schKey = schKey
        self.hasInfo = True

    def getInfo(self):
        """Get scheduler info."""
        return (self.schName, self.schKey)


def schGetTaskThreadFunc(taskManager, siBuf, schClientSock, schClientName):
    """Thread function to send task to scheduler."""
    schip, schport = schClientName
    taskInfo = taskManager.getTask(schip, schport, siBuf.getInfo()[0])
    if taskInfo:
        taskId, taskType, taskParams = taskInfo
        sendData = 'TASK %d %s %s\n' % (taskId, taskType, taskParams)
    else:
        sendData = 'TASK\n'
    try:
        schClientSock.sendall(sendData)
    except Exception, e:
        # just pass, the task may lost, but can be retried
        print schClientName, '[SCH]', e
        pass


def schCommandHandler(cmdInfo, taskManager, siBuf,
                      schClientSock, schClientName):
    """Handle the application's commands."""
    cmd, params = cmdInfo

    print schClientName, '[SCH]', cmd, 
    if params:
        print params[:40]
    else:
        print

    if cmd == 'GETTASK':
        # start a new thread, or this thread may block to
        # and wait util the task is available
        t = threading.Thread(target=schGetTaskThreadFunc,
                   args=(taskManager, siBuf, schClientSock, schClientName))
        t.setDaemon(True)
        t.start()
    elif cmd == 'RESULT':
        try:
            ps = params.split(None, 1)
            try:
                taskId, result = ps
            except:
                taskId = ps[0]
                result = ''
            schIp, schPort = schClientName
            taskManager.writeResult(taskId, schIp, schPort, result)
        except:
            print schClientName, '[SCH] invalid RESULT', params
    elif cmd == 'STATUS':
        try:
            schIp, schPort = schClientName
            taskId, statusType, statusInfo = params.split(None, 2)
            if statusType == 'nodm':
                statusInfo = statusInfo.split(None, 3)
                nodmIp, nodmPort, nodmAst, nodmRtt = statusInfo
                taskManager.updateNodeManagerStatus(taskId, schIp, schPort,
                              nodmIp, nodmPort, nodmAst, nodmRtt)
            elif statusType == 'nod':
                statusInfo = statusInfo.split(None, 4)
                nodName, nodIp, nodPort, nodAst, nodRtt = statusInfo
                taskManager.updateNodeStatus(taskId, schIp, schPort,
                              nodName, nodIp, nodPort, nodAst, nodRtt)
        except Exception, e:
            print schClientName, '[SCH]', sys._getframe().f_lineno, e
    elif cmd == 'PASSPORT':
        params = params.split(None, 1)
        if len(params) == 1:
            schName = params[0]
            schKey = ''
        else:
            schName, schKey = params
        siBuf.setInfo(schName, schKey)
        # todo: authorization


def schRecvData(taskManager, siBuf, schClientSock, schClientName, data):
    """Process data send by application."""
    siBuf.addData(data)

    while True:
        cmdInfo = siBuf.getCommand()
        if cmdInfo:
            schCommandHandler(cmdInfo, taskManager, siBuf,
                              schClientSock, schClientName)
        else:
            break


def schRecvThreadFunc(taskManager, schClientSock, schClientSocks):
    """Thread function for each scheduler connection to process data."""
    schClientName = schClientSock.getpeername()
    siBuf = SchIfBuffer()

    while True:
#        try:
            data = schClientSock.recv(CFG_SCH_BUFSIZ)
            if data == '': break
            schRecvData(taskManager, siBuf, schClientSock, schClientName, data)
#        except Exception, e:
#            print schClientName, '[SCH]', sys._getframe().f_lineno, e
    
    print schClientName, '[SCH] quit'
    schClientSock.close()
    schClientSocks.remove(schClientSock)


def startSchedulerIf(taskManager, schServerSock, schClientSocks):
    """Start interface for scheduler."""
    schRecvThreads = []

    while True:
        schClientSock, schClientAddr = schServerSock.accept()
        schClientSocks.append(schClientSock)
        print schClientAddr, '[SCH] connected'
        t = threading.Thread(target=schRecvThreadFunc,
                   args=(taskManager, schClientSock, schClientSocks))
        t.setDaemon(True)
        schRecvThreads.append(t)
        t.start()

    print 'Waiting for scheduler clients to release...',
    for t in schRecvThreads:
        t.join()
    print 'ok'

